from rich.tree import Tree
from rich.console import Console
import os

# 要排除的目录和文件列表
exclude = {
    ".git",
    "__pycache__",
    ".venv",
    "detensor_cli.egg-info",
    ".gitignore",
    "build",
    "dist",
    "tests",
    "show_tree.py",
}


def build_tree(directory, parent, max_depth, current_depth=0):
    # 如果当前深度超过最大深度，停止递归
    if current_depth >= max_depth:
        return

    for item in os.listdir(directory):
        # 如果目录或文件在排除列表中，跳过
        if item in exclude:
            continue

        full_path = os.path.join(directory, item)
        if os.path.isdir(full_path):
            branch = parent.add(f"📂 {item}")
            build_tree(full_path, branch, max_depth, current_depth + 1)  # 递归调用
        else:
            parent.add(f"📄 {item}")


console = Console()
tree = Tree("📁 Project Root")

# 设置最大深度，例如 3
max_depth = 5
build_tree(".", tree, max_depth)
console.print(tree)
